package com.fanshuai.tree.trie;

import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.lang3.StringUtils;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 多模关键词匹配前缀树
 * 多模关键词，及多个关键词，以特定符号相隔。
 * 如毛泽东&周恩来，表示同时命中"毛泽东"和"周恩来"。两个词之间最大相差M个字符
 */
public class MultiWordTrie implements ITrie {
    /**
     * 关键词之间最大相差字符
     */
    private static final int M = 6;
    /**
     * 多个关键词的分隔符
     */
    private static final char separator = '&';
    private static final String separatorString = "&";

    class Node {
        /**
         * 是否结束
         */
        public boolean isEnd;
        /**
         * 当前节点是否为分隔符
         */
        public boolean isPattern;
        /**
         * 子节点
         */
        public Map<Character, Node> subnodes;

        public Node() {
            isEnd = false;
            this.isPattern = false;
            subnodes = new HashMap<>();
        }

        public Node(boolean isPattern) {
            isEnd = false;
            this.isPattern = isPattern;
            subnodes = new HashMap<>();
        }

        public void addSubNode(char ch, Node node) {
            subnodes.put(ch, node);
        }
        public void delSubNode(char ch) {
            subnodes.remove(ch);
        }

        /**
         * 初始化前缀树时使用
         * @param ch
         * @return
         */
        public Node getSubNode(char ch) {
            return subnodes.get(ch);
        }

        /**
         * 文本匹配时使用
         * 先获取普通子节点，如果不存在，获取分隔符节点，分隔符节点不存在，返回空
         * 如果分隔符节点存在，进入多模匹配模式：
         * 顺序往后最多匹配M个字符，如果未匹配，则返回空；
         * 如果匹配，则退出多模匹配模式，继续匹配
         * @param ch
         * @return
         */
        public Node getSubNodeForMatch(char ch) {
            Node sub = subnodes.get(ch);
            if (null == sub) {
                sub = subnodes.get(separator);
            }
            return sub;
        }
    }

    private Node root;
    private int size;

    public MultiWordTrie() {
        root = new Node(false);
        size = 0;
    }

    private String trimPattern(String word) {
        int i = 0;
        int j = word.length() - 1;

        while (word.charAt(i) == '&') {
            i++;
        }

        while (word.charAt(j) == '&') {
            j--;
        }
        return word.substring(i, j + 1);
    }

    @Override
    public void addWord(String word) {
        if (StringUtils.isBlank(word)) {
            return;
        }

        //去掉开头和结束的&
        word = trimPattern(word);
        Node node = root;
        for (int i = 0; i < word.length(); i++) {
            char ch = word.charAt(i);
            if (Character.isWhitespace(ch)) {
                continue;
            }
            //连续多个分隔符时，只记录一个分隔符
            if (i > 0 && ch == separator && word.charAt(i - 1) == separator) {
                continue;
            }

            Node sub = node.getSubNode(ch);
            if (null == sub) {
                sub = new Node(ch == separator);//判断是否为分隔符节点
                node.addSubNode(ch, sub);
            }
            node = sub;

            if (i == word.length() - 1) {
                node.isEnd = true;
                size++;
            }
        }
    }

    @Override
    public void addWords(List<String> words) {
        if (CollectionUtils.isEmpty(words)) {
            return;
        }

        words.forEach(this::addWord);
    }

    @Override
    public List<String> getAllWords() {
        List<String> allWords = new ArrayList<>(size);
        StringBuilder b = new StringBuilder();

        fillWords(root, b, allWords);
        return allWords;
    }

    private void fillWords(Node node, StringBuilder b, List<String> words) {
        if (node.isEnd) {//完整关键词
            words.add(b.toString());
        }
        if (node.subnodes.size() == 0) {//叶节点
            return;
        }

        for (Character ch : node.subnodes.keySet()) {
            b.append(ch);
            fillWords(node.getSubNode(ch), b, words);

            int lastIndex = b.toString().length() - 1;
            b.deleteCharAt(lastIndex);
        }
    }

    /**
     * 匹配过程：
     * 从根节点开始匹配。如果下个字符没有匹配子节点，分2中情况：
     * 1 当前节点有没有分隔符节点，则从文本的下个字符开始重新匹配
     * 2 当前节点有分隔符节点，从分隔符节点开始，进入分隔符模式：
     *   1 下一个字符匹配成功，则退出分隔符模式，继续向下匹配
     *   2 下一个字符未匹配成功，向后遍历。如果一直未匹配，且向后遍历>=M个字符，则多模词匹配失败
     *   3 回退到文本的下一个字符，继续非分隔符模式的匹配
     * @param text
     * @return
     */
    @Override
    public List<WordToken> match(String text) {
        if (StringUtils.isBlank(text)) {
            return new ArrayList<>();
        }

        List<WordToken> tokens = new ArrayList<>();
        text = text.replace(separatorString, "");

        int pos = 0;
        int begin = 0;
        Node node = root;
        while (pos < text.length()) {
            char ch = text.charAt(pos);
            if (Character.isWhitespace(ch)) {
                pos++;
                continue;
            }

            node = node.getSubNodeForMatch(ch);
            if (null == node) {
                begin++;
                pos = begin;
                node = root;
            } else {
                if (node.isEnd) { //完全匹配
                    String word = text.substring(begin, pos + 1);
                    WordToken token = new WordToken(word, begin, pos + 1);
                    tokens.add(token);

                    pos++;
                    begin = pos;
                    node = root;
                } else {
                    if (node.isPattern) { //分隔符模式，往后遍历M个字符或匹配成功
                        Node next = node.getSubNode(ch);

                        int num = 0;
                        while (null == next && num < M) {
                            pos++;
                            if (pos >= text.length()) {
                                break;
                            }
                            ch = text.charAt(pos);
                            if (Character.isWhitespace(ch)) {
                                continue;
                            }
                            next = node.getSubNode(ch);

                            num++;
                        }

                        if (null == next) { //往后移M个字符，还未匹配成功，退出分隔符模式
                            begin++;
                            pos = begin;
                            node = root;
                        } else { //匹配到分隔符后面的词。继续往后匹配
                            node = next;
                            pos++;
                        }
                    } else {
                        pos++;
                    }
                }
            }
        }

        return tokens;
    }

    @Override
    public int size() {
        return size;
    }

    @Override
    public boolean isEmpty() {
        return size == 0;
    }
}
